import argparse
import time
import numpy as np
import os
from enum import Enum, auto
from PIL import Image
from typing import Any, Dict, List, Optional, Tuple
from matplotlib import pyplot as plt

# 设置绘图字体为黑体
plt.rcParams["font.family"] = "SimHei"


# 定义噪声类型枚举
class NoiseType(Enum):
    GAUSSIAN = auto()  # 高斯噪声
    SALT_PEPPER = auto()  # 椒盐噪声
    PERIODIC = auto()  # 周期性噪声


# 定义滤波器类型枚举
class FilterType(Enum):
    MEAN = auto()  # 均值滤波
    MEDIAN = auto()  # 中值滤波
    CONTRAHARMONIC = auto()  # 反调和平均
    BAND_REJECT = auto()  # 带阻滤波


# 噪声参数配置
NOISE_CONFIG = {
    NoiseType.GAUSSIAN: {"mean": 0, "sigma": 25},
    NoiseType.SALT_PEPPER: {"prob": 0.05},
    NoiseType.PERIODIC: {"amplitude": 50, "frequency": 0.1},
}


def add_noise(image: np.ndarray, noise_type: NoiseType) -> np.ndarray:
    """
    为图像添加指定类型的噪声
    """
    noisy = image.copy().astype(np.float32)
    h, w = image.shape

    if noise_type == NoiseType.GAUSSIAN:
        gauss = np.random.normal(
            NOISE_CONFIG[NoiseType.GAUSSIAN]["mean"],
            NOISE_CONFIG[NoiseType.GAUSSIAN]["sigma"],
            (h, w),
        )
        noisy += gauss

    elif noise_type == NoiseType.SALT_PEPPER:
        prob = NOISE_CONFIG[NoiseType.SALT_PEPPER]["prob"]
        # 生成随机矩阵
        random_matrix = np.random.rand(h, w)
        # 椒噪声（黑色像素）
        noisy[random_matrix < prob / 2] = 0
        # 盐噪声（白色像素）
        noisy[random_matrix > 1 - prob / 2] = 255

    elif noise_type == NoiseType.PERIODIC:
        x = np.arange(w)
        y = np.arange(h)
        xx, yy = np.meshgrid(x, y)
        noise = NOISE_CONFIG[NoiseType.PERIODIC]["amplitude"] * np.sin(
            2 * np.pi * NOISE_CONFIG[NoiseType.PERIODIC]["frequency"] * xx
        )
        noisy += noise

    # 限制像素值在合法范围
    return np.clip(noisy, 0, 255).astype(np.uint8)


def mean_filter(image: np.ndarray, kernel_size: int) -> Tuple[np.ndarray, float]:
    """
    均值滤波器实现
    """
    start = time.time()
    pad = kernel_size // 2
    # 添加边界填充
    padded = np.pad(image, pad, mode="edge")
    filtered = np.zeros_like(image)

    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            neighborhood = padded[i : i + kernel_size, j : j + kernel_size]
            filtered[i, j] = np.mean(neighborhood)

    return filtered, time.time() - start


def median_filter(image: np.ndarray, kernel_size: int) -> Tuple[np.ndarray, float]:
    """
    中值滤波器实现
    """
    start = time.time()
    pad = kernel_size // 2
    padded = np.pad(image, pad, mode="edge")
    filtered = np.zeros_like(image)

    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            neighborhood = padded[i : i + kernel_size, j : j + kernel_size]
            filtered[i, j] = np.median(neighborhood)

    return filtered, time.time() - start


def contraharmonic_filter(
    image: np.ndarray, kernel_size: int, Q: float
) -> Tuple[np.ndarray, float]:
    """
    反调和平均滤波器实现
    """
    start = time.time()
    pad = kernel_size // 2
    padded = np.pad(image, pad, mode="edge").astype(np.float32)
    filtered = np.zeros_like(image, dtype=np.float32)

    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            neighborhood = padded[i : i + kernel_size, j : j + kernel_size]
            numerator = np.sum(neighborhood ** (Q + 1))
            denominator = np.sum(neighborhood**Q)
            filtered[i, j] = numerator / (denominator + 1e-6)  # 防止除以零

    return np.clip(filtered, 0, 255).astype(np.uint8), time.time() - start


def band_reject_filter(
    image: np.ndarray, cutoff: float, bandwidth: float
) -> Tuple[np.ndarray, float]:
    """
    频率域带阻滤波器实现
    """
    start = time.time()
    # 傅里叶变换
    f = np.fft.fft2(image)
    fshift = np.fft.fftshift(f)

    # 创建滤波器
    rows, cols = image.shape
    crow, ccol = rows // 2, cols // 2
    x = np.linspace(-0.5, 0.5, cols)
    y = np.linspace(-0.5, 0.5, rows)
    xx, yy = np.meshgrid(x, y)
    d = np.sqrt(xx**2 + yy**2)

    # 高斯带阻滤波器
    mask = 1 - np.exp(-((d - cutoff) ** 2 / (2 * (bandwidth / 2) ** 2)))

    # 应用滤波器
    fshift_filtered = fshift * mask

    # 逆变换
    f_ishift = np.fft.ifftshift(fshift_filtered)
    filtered = np.fft.ifft2(f_ishift)
    result = np.abs(filtered)

    return np.clip(result, 0, 255).astype(np.uint8), time.time() - start


def plot_results(
    original: np.ndarray,
    noisy: np.ndarray,
    filtered: np.ndarray,
    title: str,
    save_path: Optional[str],
):
    """
    绘制并保存对比图
    """
    plt.figure(figsize=(15, 5))
    plt.subplot(131)
    plt.imshow(original, cmap="gray")
    plt.title("原始图像")
    plt.axis("off")
    plt.subplot(132)
    plt.imshow(noisy, cmap="gray")
    plt.title("噪声图像")
    plt.axis("off")
    plt.subplot(133)
    plt.imshow(filtered, cmap="gray")
    plt.title(title)
    plt.axis("off")
    if save_path:
        plt.savefig(save_path)
    plt.close()


def main():
    parser = argparse.ArgumentParser(description="图像去噪实验")
    parser.add_argument("--image_path", help="输入图像路径")
    parser.add_argument("--output_dir", help="输出目录", default=None)
    args = parser.parse_args()

    if args.image_path.startswith('"') and args.image_path.endswith('"'):
        args.image_path = args.image_path[1:-1]

    # 读取图像并转换为灰度
    original = np.array(Image.open(args.image_path).convert("L"))

    # 实验流程
    for noise_type in NoiseType:
        # 添加噪声
        noisy = add_noise(original, noise_type)

        # 应用不同滤波器
        filters: List[Tuple[FilterType, Dict[str, Any]]] = [
            (FilterType.MEAN, {"kernel_size": 3}),
            (FilterType.MEDIAN, {"kernel_size": 3}),
            (FilterType.CONTRAHARMONIC, {"kernel_size": 3, "Q": 1.5}),
            (FilterType.BAND_REJECT, {"cutoff": 0.2, "bandwidth": 0.1}),
        ]

        for filter_type, params in filters:
            # 执行滤波
            if filter_type == FilterType.MEAN:
                filtered, time_used = mean_filter(noisy, **params)
            elif filter_type == FilterType.MEDIAN:
                filtered, time_used = median_filter(noisy, **params)
            elif filter_type == FilterType.CONTRAHARMONIC:
                filtered, time_used = contraharmonic_filter(noisy, **params)
            else:
                filtered, time_used = band_reject_filter(noisy, **params)

            # 保存结果
            title = f"{filter_type.name} 处理 ({time_used:.2f}s)"
            filename = (
                os.path.join(
                    args.output_dir, f"{noise_type.name}_{filter_type.name}.png"
                )
                if args.output_dir
                else None
            )
            plot_results(original, noisy, filtered, title, filename)


if __name__ == "__main__":
    main()
